from numpy import sin, linspace, pi, log10, where, copy, power
from numpy.fft import fftfreq
from scipy import fft, arange, ifft
import numpy

import matplotlib
from scipy.io import wavfile
import scipy.io.wavfile
matplotlib.use('WXAgg')
from fractions import Fraction

from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg as FigureCanvas
from matplotlib.backends.backend_wx import NavigationToolbar2Wx
from matplotlib.figure import Figure

import wx

class Equaliser(wx.Panel):

    def __init__(self, parent, model):
        wx.Panel.__init__(self, parent=parent, id=wx.ID_ANY)
        self.parent = parent
        self.model = model

        self.figureSample = Figure()
        self.axesSample = self.figureSample.add_subplot(111)
        self.canvasSample = FigureCanvas(self, wx.ID_ANY, self.figureSample)
        self.canvasSample.SetSize((500, 400))

        self.figureFreq = Figure()
        self.axesFreq = self.figureFreq.add_subplot(111)
        self.canvasFreq = FigureCanvas(self, wx.ID_ANY, self.figureFreq)
        self.canvasFreq.SetSize((500, 400))

        
        self.option = wx.ComboBox(self, wx.ID_ANY, pos=(50, 170), size=(150, -1), choices=[], style=wx.CB_READONLY)
        self.option.Bind(wx.EVT_COMBOBOX, self.OnSelect)

        self.firstBand = wx.Slider(self, wx.ID_ANY, 1000, 0, 2000, wx.DefaultPosition, (50, -1),
                              wx.SL_AUTOTICKS | wx.SL_VERTICAL)
        self.firstBand.Bind(wx.EVT_SLIDER, self.OnAdjustFirst)
        self.secondBand = wx.Slider(self, wx.ID_ANY, 1000, 0, 2000, wx.DefaultPosition, (50, -1),
                                     wx.SL_AUTOTICKS | wx.SL_VERTICAL)
        self.secondBand.Bind(wx.EVT_SLIDER, self.OnAdjustSecond)
        self.thirdBand = wx.Slider(self, wx.ID_ANY, 1000, 0, 2000, wx.DefaultPosition, (50, -1),
                              wx.SL_AUTOTICKS | wx.SL_VERTICAL)
        self.thirdBand.Bind(wx.EVT_SLIDER, self.OnAdjustThird)
        self.fourthBand = wx.Slider(self, wx.ID_ANY, 1000, 0, 2000, wx.DefaultPosition, (50, -1),
                                     wx.SL_AUTOTICKS | wx.SL_VERTICAL)
        self.fourthBand.Bind(wx.EVT_SLIDER, self.OnAdjustFourth)
        self.fifthBand = wx.Slider(self, wx.ID_ANY, 1000, 0, 2000, wx.DefaultPosition, (50, -1),
                              wx.SL_AUTOTICKS | wx.SL_VERTICAL)
        self.fifthBand.Bind(wx.EVT_SLIDER, self.OnAdjustFifth)
        self.sixthBand = wx.Slider(self, wx.ID_ANY, 1000, 0, 2000, wx.DefaultPosition, (50, -1),
                                     wx.SL_AUTOTICKS | wx.SL_VERTICAL)
        self.sixthBand.Bind(wx.EVT_SLIDER, self.OnAdjustSixth)

        save = wx.Button(self, wx.ID_ANY, "Save")
        save.Bind(wx.EVT_BUTTON, self.OnSave)

        sizer = wx.BoxSizer(wx.VERTICAL)
        graphs = wx.BoxSizer(wx.HORIZONTAL)
        equaliser = wx.BoxSizer(wx.HORIZONTAL)
        sizer.Add(self.option, 0, wx.ALL, 5)
        sizer.Add(graphs, 0, wx.ALL, 5)
        sizer.Add(equaliser, 0, wx.ALL, 5)

        slider1 = wx.BoxSizer(wx.VERTICAL)
        slider2 = wx.BoxSizer(wx.VERTICAL)
        slider3 = wx.BoxSizer(wx.VERTICAL)
        slider4 = wx.BoxSizer(wx.VERTICAL)
        slider5 = wx.BoxSizer(wx.VERTICAL)
        slider6 = wx.BoxSizer(wx.VERTICAL)

        band1 = wx.StaticText(self, label="4kHz")
        band2 = wx.StaticText(self, label="8kHz")
        band3 = wx.StaticText(self, label="12kHz")
        band4 = wx.StaticText(self, label="16kHz")
        band5 = wx.StaticText(self, label="20kHz")
        band6 = wx.StaticText(self, label="24kHz")

        slider1.Add(band1)
        slider1.Add(self.firstBand)
        
        slider2.Add(band2)
        slider2.Add(self.secondBand)

        slider3.Add(band3)
        slider3.Add(self.thirdBand)

        slider4.Add(band4)
        slider4.Add(self.fourthBand)

        slider5.Add(band5)
        slider5.Add(self.fifthBand)

        slider6.Add(band6)
        slider6.Add(self.sixthBand)

        equaliser.Add(slider1, 0, wx.ALL, 5)
        equaliser.Add(slider2, 0, wx.ALL, 5)
        equaliser.Add(slider3, 0, wx.ALL, 5)
        equaliser.Add(slider4, 0, wx.ALL, 5)
        equaliser.Add(slider5, 0, wx.ALL, 5)
        equaliser.Add(slider6, 0, wx.ALL, 5)

        graphs.Add(self.canvasSample, 0, wx.ALL, 5)
        graphs.Add(self.canvasFreq, 0, wx.ALL, 5)

        sizer.Add(save, 0, wx.ALL, 5)

        self.SetSizer(sizer)

        self.freqWorking = None

    def OnSelect(self, event):
        self.axesSample.clear()
        self.axesFreq.clear()
        filename = self.option.GetValue()
        location = self.model.getFile(filename)
        self.audio = self.model.getAudio(location)
        x = numpy.array(range(0, len(self.audio[1])))/44100.0
        self.plotSample = self.axesSample.plot(x, self.audio[1]/(2.0**15))
        self.figureSample.canvas.draw()
        n = len(self.audio[1])
        k = arange(n)
        T= n/self.audio[0]
        self.freqx = k/T
        self.freqx = self.freqx[:n/2]
        self.freqy = fft(self.audio[1])/n
        self.freqy = self.freqy[:n/2]
        self.freqy = 20*log10(self.freqy)
        self.freqy[where(self.freqy<0)] = 0
        self.plotFreq = self.axesFreq.plot(self.freqx, self.freqy, 'g')
        self.figureFreq.canvas.draw()
        self.length = len(self.audio[1])
        self.freqWorking = copy(self.freqy)
        self.output = self.audio[1]

    def invert(self, samples):
        self.output = ifft(power(samples/20, 10)*len(self.audio[1]))
        self.axesSample.clear()
        x = numpy.array(range(0, len(self.output)))/44100.0
        self.plotSample = self.axesSample.plot(x, self.output/(2.0**15))
        self.figureSample.canvas.draw()

    def OnAdjustFirst(self, event):
        gain = self.firstBand.GetValue()/1000.0
        self.freqWorking = copy(self.freqy)
        for sample in range(4000):
            self.freqWorking[sample] = self.freqWorking[sample] * gain
        self.axesFreq.clear()
        self.axesFreq.plot(self.freqx, self.freqWorking, 'g')
        self.figureFreq.canvas.draw()
        self.invert(self.freqWorking)

    def OnAdjustSecond(self, event):
        gain = self.firstBand.GetValue()/1000.0
        self.freqWorking = copy(self.freqy)
        for sample in range(4000,8000):
            self.freqWorking[sample] = self.freqWorking[sample] * gain
        self.axesFreq.clear()
        self.axesFreq.plot(self.freqx, self.freqWorking, 'g')
        self.figureFreq.canvas.draw()

    def OnAdjustThird(self, event):
        gain = self.firstBand.GetValue()/1000.0
        self.freqWorking = copy(self.freqy)
        for sample in range(8000,12000):
            self.freqWorking[sample] = self.freqWorking[sample] * gain
        self.axesFreq.clear()
        self.axesFreq.plot(self.freqx, self.freqWorking, 'g')
        self.figureFreq.canvas.draw()

    def OnAdjustFourth(self, event):
        gain = self.firstBand.GetValue()/1000.0
        self.freqWorking = copy(self.freqy)
        for sample in range(12000,16000):
            self.freqWorking[sample] = self.freqWorking[sample] * gain
        self.axesFreq.clear()
        self.axesFreq.plot(self.freqx, self.freqWorking, 'g')
        self.figureFreq.canvas.draw()

    def OnAdjustFifth(self, event):
        gain = self.firstBand.GetValue()/1000.0
        self.freqWorking = copy(self.freqy)
        for sample in range(16000,20000):
            self.freqWorking[sample] = self.freqWorking[sample] * gain
        self.axesFreq.clear()
        self.axesFreq.plot(self.freqx, self.freqWorking, 'g')
        self.figureFreq.canvas.draw()

    def OnAdjustSixth(self, event):
        gain = self.firstBand.GetValue()/1000.0
        self.freqWorking = copy(self.freqy)
        for sample in range(20000,24000):
            self.freqWorking[sample] = self.freqWorking[sample] * gain
        self.axesFreq.clear()
        self.axesFreq.plot(self.freqx, self.freqWorking, 'g')
        self.figureFreq.canvas.draw()

    def OnSave(self, event):
        dlg = wx.FileDialog(
            self, message="Choose a file",
            defaultFile="",
            style=wx.SAVE | wx.CHANGE_DIR
            )
        if dlg.ShowModal() == wx.ID_OK:
            path = dlg.GetPath()
            wavfile.write(path, self.audio[0], self.output.astype(self.audio[1].dtype))

    def update(self):
        self.option.Clear()
        filenames = self.model.getFilenames()
        for filename in filenames:
            self.option.Append(filename.keys()[0])
